#include "downloader_logic.h"
#include "downloader_window.h" // Included again for UIData definition access
#include <curl/curl.h>
#include <sstream>
#include <iostream> // Debug output
#include <cstdio>   // For std::remove
#include <cmath>    // For std::ceil
#include <algorithm> // For std::min

// Assumed block size if file size is unknown (less reliable)
const long long ASSUMED_BLOCK_SIZE_LOGIC = 10 * 1024 * 1024; // 10MB

DownloaderLogic::DownloaderLogic(std::string url, std::string savePath, int numThreads,
                                 std::atomic<bool>& cancelFlag,
                                 Glib::Dispatcher& dispatcher,
                                 DownloaderWindow::UIData& uiData,
                                 std::mutex& uiDataMutex)
    : m_url(std::move(url)),
      m_savePath(std::move(savePath)),
      m_numThreads(std::max(1, numThreads)), // Ensure at least 1 thread
      m_totalFileSize(-1),
      m_totalDownloadedSize(0),
      m_active_workers(0),
      m_worker_error_occurred(false),
      m_cancel_flag_ref(cancelFlag),
      m_dispatcher_ref(dispatcher),
      m_ui_data_ref(uiData),
      m_ui_data_mutex_ref(uiDataMutex)
{
    // Derive base temp filename from save path
    size_t last_slash = m_savePath.find_last_of("/\\");
    std::string path_part = (last_slash == std::string::npos) ? "" : m_savePath.substr(0, last_slash + 1);
    std::string name_part = (last_slash == std::string::npos) ? m_savePath : m_savePath.substr(last_slash + 1);
    m_baseTempFilename = path_part + "." + name_part + ".part"; // e.g., /path/.file.txt.part
    std::cout << "Logic created. Base temp name: " << m_baseTempFilename << std::endl;
}

DownloaderLogic::~DownloaderLogic() {
     std::cout << "DownloaderLogic destructor." << std::endl;
     // Destructor might be called if thread creation fails or window is closed early.
     // Ensure workers are joined if they were started (though run_download should handle join)
     for (auto& th : m_worker_threads) {
         if (th.joinable()) {
             th.join();
         }
     }
     cleanup_temp_files(); // Attempt cleanup just in case
}


void DownloaderLogic::run_download() {
    std::cout << "run_download started on thread " << std::this_thread::get_id() << std::endl;

    report_status("Getting file size...");
    m_totalFileSize = get_file_size_internal();

    if (m_cancel_flag_ref) {
        report_status("Cancelled", false, true); // Mark as finished (cancelled)
        cleanup_temp_files();
        return;
    }

    if (m_totalFileSize == 0) {
        report_status("File size is 0 bytes. Creating empty file.", false, true);
        // Create an empty file
        std::ofstream empty_file(m_savePath, std::ios::binary);
        empty_file.close();
        report_final_progress(); // Show 100% for 0-byte file
        return;
    }

     if (m_totalFileSize < 0) {
          report_status("Warning: Could not determine file size. Attempting download with 1 thread.", false, false);
          m_numThreads = 1; // Fallback to single thread if size is unknown
     } else {
          report_status(std::string("File size: ") + std::to_string(m_totalFileSize / (1024*1024)) + " MB. Starting download...", false, false);
     }

    prepare_download_blocks();

    if (m_tempFilePaths.empty()) {
        report_status("Failed to prepare download blocks.", true, true, "Block preparation resulted in no tasks.");
        return;
    }

    launch_worker_threads();

    // Wait for all worker threads to complete
    std::cout << "Main logic thread waiting for " << m_worker_threads.size() << " workers..." << std::endl;
    for (auto& th : m_worker_threads) {
        if (th.joinable()) {
            th.join();
        }
    }
     std::cout << "All worker threads finished." << std::endl;
     report_final_progress(); // Ensure total progress is reported accurately


    // Check for cancellation or worker errors *after* joining
    if (m_cancel_flag_ref) {
        report_status("Cancelled", false, true);
        cleanup_temp_files();
        return;
    }
    if (m_worker_error_occurred) {
        report_status("Download failed due to worker error.", true, true, "One or more download blocks failed.");
        cleanup_temp_files();
        return;
    }

    // Merge files if everything succeeded so far
    report_status("Merging downloaded parts...");
    if (merge_files_internal()) {
        report_status("Download complete!", false, true);
        cleanup_temp_files(); // Clean up on success
    } else {
        // Merge failed (error reported inside merge_files_internal)
        // Cleanup should happen inside merge if it fails partially
    }
     std::cout << "run_download finished on thread " << std::this_thread::get_id() << std::endl;
}

void DownloaderLogic::cleanup_temp_files() {
     std::cout << "Cleaning up temporary files..." << std::endl;
     for (const auto& temp_file : m_tempFilePaths) {
          if (std::remove(temp_file.c_str()) == 0) {
               std::cout << "Removed: " << temp_file << std::endl;
          } else {
               // Non-fatal, maybe file didn't exist or permissions issue
               std::cerr << "Warning: Could not remove temp file: " << temp_file << std::endl;
          }
     }
     m_tempFilePaths.clear(); // Clear the list
}

long long DownloaderLogic::get_file_size_internal() {
    CURL *curl = curl_easy_init();
    long long file_size = -1;
    if (!curl) {
        report_status("Failed to init curl for size check.", true, true, "curl_easy_init returned null");
        return -1;
    }

    // Progress data for cancellation check
    CurlWorkerData progress_data;
    progress_data.logic_ptr = this;

    curl_easy_setopt(curl, CURLOPT_URL, m_url.c_str());
    curl_easy_setopt(curl, CURLOPT_NOBODY, 1L); // HEAD request
    curl_easy_setopt(curl, CURLOPT_FOLLOWLOCATION, 1L);
    // Enable progress function to check for cancellation
    curl_easy_setopt(curl, CURLOPT_NOPROGRESS, 0L);
    curl_easy_setopt(curl, CURLOPT_PROGRESSFUNCTION, progress_callback_static);
    curl_easy_setopt(curl, CURLOPT_PROGRESSDATA, &progress_data);

    CURLcode res = curl_easy_perform(curl);

    if (m_cancel_flag_ref) {
         std::cout << "File size check cancelled by user." << std::endl;
         // Status already set by run_download or will be set soon
         file_size = -2; // Indicate cancellation occurred
    } else if (res == CURLE_OK) {
        double content_length = -1.0;
        res = curl_easy_getinfo(curl, CURLINFO_CONTENT_LENGTH_DOWNLOAD, &content_length);
        if (res == CURLE_OK && content_length >= 0.0) {
            file_size = static_cast<long long>(content_length);
            std::cout << "Retrieved file size: " << file_size << " bytes." << std::endl;
        } else {
            std::cerr << "Could not get content length from HEAD request: " << curl_easy_strerror(res) << std::endl;
            file_size = -1; // Indicate size unknown
        }
    } else {
        std::cerr << "HEAD request failed: " << curl_easy_strerror(res) << std::endl;
        file_size = -1; // Indicate size unknown or error
    }

    curl_easy_cleanup(curl);
    return file_size;
}

void DownloaderLogic::prepare_download_blocks() {
    m_tempFilePaths.clear();
    long long current_pos = 0;

    if (m_totalFileSize > 0) {
         // Calculate block size, ensuring it's at least 1 byte and handles division correctly
         long long block_size = std::max(1LL, m_totalFileSize / m_numThreads);
         // Adjust block size slightly to distribute remainder (optional, last block handles it too)
         // block_size = static_cast<long long>(std::ceil(static_cast<double>(m_totalFileSize) / m_numThreads));

        for (int i = 0; i < m_numThreads; ++i) {
            DownloadBlock block;
            block.block_id = i;
            block.url = m_url;
            block.start = current_pos;
             // Calculate end: current + size - 1, but not exceeding total size
             // Also handle the last block explicitly
            block.end = (i == m_numThreads - 1) ? m_totalFileSize - 1 : current_pos + block_size - 1;
             // Ensure end doesn't go past total size if block size calculation was imperfect
            block.end = std::min(block.end, m_totalFileSize - 1);

            // Ensure start <= end (handles tiny files where block_size might be > remaining)
            if (block.start > block.end) {
                 if (i == m_numThreads -1 && block.start == m_totalFileSize) {
                     // This thread gets a 0-byte block if division was perfect, skip it? Or handle 0-byte download?
                      std::cout << "Skipping 0-byte block for thread " << i << std::endl;
                      continue; // Don't create a task for a 0-byte range
                 } else {
                      // Should not happen with std::min, but as safety:
                       report_status("Block calculation error.", true, true, "Start position exceeded end position.");
                       m_tempFilePaths.clear(); // Prevent starting download
                       return;
                 }
            }

            block.temp_filename = m_baseTempFilename + std::to_string(i);
            m_tempFilePaths.push_back(block.temp_filename);

            m_worker_threads.emplace_back(&DownloaderLogic::download_block_worker, this, block);
            current_pos = block.end + 1; // Start of next block

             // Safety check: if current_pos reaches end, stop creating blocks
             if (current_pos >= m_totalFileSize) break;
        }
    } else { // File size unknown or error (-1)
        // Fallback to single block/thread (m_numThreads was set to 1 earlier)
        DownloadBlock block;
        block.block_id = 0;
        block.url = m_url;
        block.start = 0;
        block.end = -1; // Indicate download until end of stream
        block.temp_filename = m_baseTempFilename + std::to_string(0);
        m_tempFilePaths.push_back(block.temp_filename);
        m_worker_threads.emplace_back(&DownloaderLogic::download_block_worker, this, block);
    }
     m_active_workers = m_worker_threads.size(); // Set initial count
}


void DownloaderLogic::launch_worker_threads() {
    // The threads were already created and stored in m_worker_threads
    // during prepare_download_blocks. This function might not be needed,
    // or could be used to detach threads if join is handled differently.
    // Current implementation joins them in run_download.
    std::cout << "Prepared " << m_worker_threads.size() << " worker threads." << std::endl;
}

// This function runs in a separate worker thread
void DownloaderLogic::download_block_worker(DownloadBlock blockInfo) {
    std::cout << "Worker " << blockInfo.block_id << ": Started. Range " << blockInfo.start << "-" << blockInfo.end << std::endl;

    CURL *curl = curl_easy_init();
    if (!curl) {
        std::cerr << "Worker " << blockInfo.block_id << ": Failed to initialize curl." << std::endl;
        m_worker_error_occurred = true;
        m_active_workers--;
        return;
    }

    // Prepare file stream
    std::ofstream outfile(blockInfo.temp_filename, std::ios::binary);
    if (!outfile.is_open()) {
        std::cerr << "Worker " << blockInfo.block_id << ": Failed to open temp file " << blockInfo.temp_filename << std::endl;
        curl_easy_cleanup(curl);
        m_worker_error_occurred = true;
        m_active_workers--;
        return;
    }

    // Prepare data for callbacks
    CurlWorkerData workerData;
    workerData.logic_ptr = this;
    workerData.outfile = &outfile;
    workerData.block_start_pos = blockInfo.start;
    workerData.block_current_pos = blockInfo.start; // Tracks position within the block/file
    workerData.block_total_size = (blockInfo.end >= blockInfo.start) ? (blockInfo.end - blockInfo.start + 1) : 0; // Size of this block


    // Set curl options
    curl_easy_setopt(curl, CURLOPT_URL, blockInfo.url.c_str());
    curl_easy_setopt(curl, CURLOPT_FOLLOWLOCATION, 1L);
    curl_easy_setopt(curl, CURLOPT_WRITEFUNCTION, write_callback_static);
    curl_easy_setopt(curl, CURLOPT_WRITEDATA, &workerData);
    curl_easy_setopt(curl, CURLOPT_NOPROGRESS, 0L);
    curl_easy_setopt(curl, CURLOPT_PROGRESSFUNCTION, progress_callback_static);
    curl_easy_setopt(curl, CURLOPT_PROGRESSDATA, &workerData); // Pass worker data to progress callback too
    // Set Range header if block has a valid range
    if (blockInfo.end >= blockInfo.start && blockInfo.start >= 0) {
        std::string range = std::to_string(blockInfo.start) + "-" + std::to_string(blockInfo.end);
        curl_easy_setopt(curl, CURLOPT_RANGE, range.c_str());
        std::cout << "Worker " << blockInfo.block_id << ": Set Range: " << range << std::endl;
    } else {
         std::cout << "Worker " << blockInfo.block_id << ": No Range header (full download or unknown size)." << std::endl;
    }
     // Optional: Add timeouts, user agent etc.
     curl_easy_setopt(curl, CURLOPT_CONNECTTIMEOUT, 20L); // 20 seconds connect timeout
     // curl_easy_setopt(curl, CURLOPT_TIMEOUT, 120L); // 120 seconds total time (careful with large files)


    // Perform the download
    CURLcode res = curl_easy_perform(curl);

    // Cleanup file stream (important!)
    outfile.close();

    // Check result
    if (m_cancel_flag_ref) {
         std::cout << "Worker " << blockInfo.block_id << ": Download cancelled by user." << std::endl;
         // Don't set worker error flag if cancelled gracefully
    } else if (res != CURLE_OK) {
        std::cerr << "Worker " << blockInfo.block_id << ": Download failed: " << curl_easy_strerror(res) << std::endl;
        m_worker_error_occurred = true;
    } else {
         // Check HTTP response code
         long http_code = 0;
         curl_easy_getinfo(curl, CURLINFO_RESPONSE_CODE, &http_code);
         // Accept 2xx (Successful) and especially 206 (Partial Content for ranges)
         if (http_code >= 200 && http_code < 300) {
              std::cout << "Worker " << blockInfo.block_id << ": Finished successfully (HTTP " << http_code << ")." << std::endl;
         } else {
              std::cerr << "Worker " << blockInfo.block_id << ": Download failed with HTTP status: " << http_code << std::endl;
              m_worker_error_occurred = true;
         }
    }

    curl_easy_cleanup(curl);
    m_active_workers--;
    std::cout << "Worker " << blockInfo.block_id << ": Exiting. Active workers remaining: " << m_active_workers << std::endl;
}

bool DownloaderLogic::merge_files_internal() {
    std::cout << "Merging files into: " << m_savePath << std::endl;
    std::ofstream final_file(m_savePath, std::ios::binary);
    if (!final_file.is_open()) {
        report_status("Merge failed: Cannot open final file.", true, true, "Failed to open " + m_savePath + " for writing.");
        return false;
    }

    bool merge_ok = true;
    for (const auto& temp_file : m_tempFilePaths) {
        if (m_cancel_flag_ref) {
             merge_ok = false; // Don't consider merge successful if cancelled
             report_status("Cancelled during merge", false, true); // Report cancellation
             break;
        }

        std::ifstream part_file(temp_file, std::ios::binary);
        if (part_file.is_open()) {
            final_file << part_file.rdbuf(); // Stream content
            part_file.close();
            if (!final_file.good()) {
                 report_status("Merge failed: Error writing to final file.", true, true, "Disk full or I/O error during merge.");
                 merge_ok = false;
                 break;
            }
        } else {
            report_status("Merge failed: Cannot open temp file.", true, true, "Failed to open " + temp_file + " for reading.");
            merge_ok = false;
            break;
        }
    }

    final_file.close(); // Close the final file

    if (!merge_ok) {
        std::cerr << "Merge process failed or was cancelled. Attempting to delete incomplete final file." << std::endl;
        std::remove(m_savePath.c_str()); // Delete potentially incomplete/corrupt file
        cleanup_temp_files(); // Clean up parts regardless
    }
    // If merge_ok is true, success/completion is reported by run_download after calling merge.
    // Cleanup of temp files also happens after successful merge in run_download.

    return merge_ok;
}


// --- Reporting Methods ---

void DownloaderLogic::report_progress_increment(long long downloaded_increment) {
    // Atomically update the total downloaded size
    m_totalDownloadedSize += downloaded_increment;

    // Update shared UI data structure (under lock)
    {
        std::lock_guard<std::mutex> lock(m_ui_data_mutex_ref);
        m_ui_data_ref.current_downloaded = m_totalDownloadedSize; // Use atomic value for accuracy
        if (m_totalFileSize > 0) {
            m_ui_data_ref.percentage = std::min(1.0, static_cast<double>(m_ui_data_ref.current_downloaded) / m_totalFileSize);
        } else {
            m_ui_data_ref.percentage = 0.0; // Or some other indicator for unknown size
        }
        // Keep status message as is unless specifically changing it
    }
    m_dispatcher_ref.emit(); // Notify GUI thread
}

void DownloaderLogic::report_final_progress() {
     // Ensure the final downloaded size and percentage are accurate
     {
          std::lock_guard<std::mutex> lock(m_ui_data_mutex_ref);
          m_ui_data_ref.current_downloaded = m_totalDownloadedSize;
          if (m_totalFileSize > 0) {
               // Handle potential floating point inaccuracies or slight overallocation
               if (m_totalDownloadedSize >= m_totalFileSize && !m_worker_error_occurred && !m_cancel_flag_ref) {
                    m_ui_data_ref.percentage = 1.0;
                    m_ui_data_ref.current_downloaded = m_totalFileSize; // Cap at total size
               } else {
                    m_ui_data_ref.percentage = std::min(1.0, static_cast<double>(m_ui_data_ref.current_downloaded) / m_totalFileSize);
               }
          } else if (m_totalFileSize == 0) {
                m_ui_data_ref.percentage = 1.0; // 0-byte file is 100% complete
          } else {
               m_ui_data_ref.percentage = 0.0; // Unknown size, indeterminate
          }
     }
     m_dispatcher_ref.emit();
}


void DownloaderLogic::report_status(const std::string& message, bool is_error, bool is_finished, const std::string& error_details) {
    std::cout << "Reporting Status: [" << (is_error?"Error":"Info") << "/" << (is_finished?"Finished":"Ongoing") << "] " << message;
    if (!error_details.empty()) std::cout << " Details: " << error_details;
    std::cout << std::endl;

    {
        std::lock_guard<std::mutex> lock(m_ui_data_mutex_ref);
        m_ui_data_ref.status_message = message;
        m_ui_data_ref.has_error = is_error;
        m_ui_data_ref.is_finished = is_finished; // Mark as finished (could be success, error, or cancel)
        m_ui_data_ref.error_message = error_details;

        // Update percentage based on final state if applicable
        if (is_finished) {
             if (!is_error && !m_cancel_flag_ref) { // Success
                  if (m_totalFileSize >= 0) m_ui_data_ref.percentage = 1.0;
             } else if (is_error) { // Error
                   // Keep current percentage or reset? Maybe keep for context.
             } else { // Cancelled
                   // Keep current percentage
             }
             // Ensure running flag is off if we mark as finished
             // m_ui_data_ref.is_running = false; // This should be done in the GUI thread handler
        }
    }
    m_dispatcher_ref.emit(); // Notify GUI thread
}


// --- Static Curl Callbacks ---

size_t DownloaderLogic::write_callback_static(char *ptr, size_t size, size_t nmemb, void *userdata) {
    CurlWorkerData* data = static_cast<CurlWorkerData*>(userdata);
    if (!data || !data->logic_ptr || !data->outfile) {
        return 0; // Indicate error
    }

    // Check for cancellation before writing
    if (data->logic_ptr->m_cancel_flag_ref) {
        return 0; // Stop download by indicating error
    }

    size_t realSize = size * nmemb;
    data->outfile->write(ptr, realSize);

    if (data->outfile->good()) {
        data->logic_ptr->report_progress_increment(realSize); // Report increment to main logic
        data->block_current_pos += realSize; // Track position within this block's data
        return realSize; // Success
    } else {
        std::cerr << "Write callback: Failed to write to temp file!" << std::endl;
        data->logic_ptr->m_worker_error_occurred = true; // Signal error
        return 0; // Indicate write error to curl
    }
}

int DownloaderLogic::progress_callback_static(void *clientp, double dltotal, double dlnow, double ultotal, double ulnow) {
    CurlWorkerData* data = static_cast<CurlWorkerData*>(clientp);
    if (data && data->logic_ptr) {
        if (data->logic_ptr->m_cancel_flag_ref) {
            // std::cout << "Progress callback: Cancel detected." << std::endl;
            return 1; // Non-zero return aborts the transfer
        }
    }
    return 0; // Zero continues the transfer
}
